package com.example.qwen02;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.embeddings.TextEmbedding;
import com.alibaba.dashscope.embeddings.TextEmbeddingParam;
import com.alibaba.dashscope.embeddings.TextEmbeddingResult;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.tools.FunctionDefinition;
import com.alibaba.dashscope.tools.ToolCallBase;
import com.alibaba.dashscope.tools.ToolCallFunction;
import com.alibaba.dashscope.utils.JsonUtils;
import com.aliyun.dashvector.DashVectorClient;
import com.aliyun.dashvector.DashVectorCollection;
import com.aliyun.dashvector.models.Doc;
import com.aliyun.dashvector.models.Vector;
import com.aliyun.dashvector.models.requests.QueryDocRequest;
import com.aliyun.dashvector.models.responses.Response;
import com.example.qwen02.Fun01;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.victools.jsonschema.generator.*;

import java.util.*;

public class Rag02c {
    public static class WeekTool{
        public  String weekToday(){
            String[] names=new String[]{"星期日","星期一","星期二","星期三","星期四","星期五","星期六"};
            Calendar c=Calendar.getInstance();
            int w=c.get(Calendar.DAY_OF_WEEK);
            return "今天是"+names[w-1];
        }
    }
    public static void main(String[] args) throws NoApiKeyException, InputRequiredException {

        Scanner scanner = new Scanner(System.in);
        System.out.println("请输入问题：");
        String question = scanner.nextLine();
        System.out.println("请输入最大数量：");
        int maxTokens = scanner.nextInt();

        SchemaGeneratorConfigBuilder configBuilder =
                new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON);
        SchemaGeneratorConfig config = configBuilder.with(Option.EXTRA_OPEN_API_FORMAT_VALUES)
                .without(Option.FLATTENED_ENUMS_FROM_TOSTRING).build();
        SchemaGenerator generator = new SchemaGenerator(config);

        ObjectNode jsonSchema_week = generator.generateSchema(Fun01.WeekTool.class);
        FunctionDefinition fd_week = FunctionDefinition.builder().name("get_current_week").description("获取今天星期几")
                .parameters(JsonUtils.parseString(jsonSchema_week.toString()).getAsJsonObject()).build();


        //从向量数据库根据用户的问题，去过滤出相关的信息段落，减少喂入大模型的tokens
        //但是向量数据库只能使用数字矩阵比对相识度，所以需要先将问题向量化，然后和数据库的向量进行比对，得到最相关的信息段落
        //把文字转换成向量
        TextEmbeddingParam param = TextEmbeddingParam
                .builder()
                .model(TextEmbedding.Models.TEXT_EMBEDDING_V1)
                .texts(Arrays.asList(question)).build();
        TextEmbedding textEmbedding = new TextEmbedding();
        TextEmbeddingResult result = textEmbedding.call(param);
        System.out.println(result);

        List<Double> embedding_double = result.getOutput().getEmbeddings().get(0).getEmbedding();
        List<Float> embedding_float = new ArrayList<>(embedding_double.size());
        for(Double d:embedding_double)
            embedding_float.add(d.floatValue());
        //插入向量数据库
        // 构建Vector
        Vector vector = Vector.builder().value(embedding_float).build();
        //查询数据(相似度查询)
        QueryDocRequest request = QueryDocRequest.builder()
                .vector(vector)
                .topk(2)
                .includeVector(true)
                .outputField("raw")
                .build();
        // 进行Doc检索
        DashVectorClient client = new DashVectorClient("sk-a1WoRTf9XCDOVlYL1u8vLKqmf4NAE6D51805C198411EFBFA83AFB039D59B0"
                , "vrs-cn-o493r4u150003y.dashvector.cn-hangzhou.aliyuncs.com");
        DashVectorCollection collection = client.get("courses");
        Response<List<Doc>> response = collection.query(request);
        StringBuffer ragtxt = new StringBuffer();
        for(Doc doc:response.getOutput()){
            String cxt=doc.getFields().get("raw").toString();
            System.out.println(doc.getScore()+"-->"+cxt);
            ragtxt.append(cxt);
            ragtxt.append("\n");
        }

        Message systemMsg = Message.builder().role(Role.SYSTEM.getValue())
                .content("角色:你是一个教务小助手;任务:请根据我们提供的课程信息，回答学生对课程安排的查询;例如:我说:软件3班,周一早上上什么课，你就回答:软件3班在周一早上上高等数学;举例完毕。\" +\n" +
                        "注意回答的要求:如果不能查询到准确信息，请回答，查询不到此信息，千万不要乱说。另外遇到问题，还可以查询工具方法。提供以下课程表资料供你参考：\\n\"+ragtxt.toString()")
                .build();
        Message userMsg = Message.builder().role(Role.USER.getValue())
                .content(question)
                .build();

        List<Message> msgList=new LinkedList<>();
        msgList.add(systemMsg);
        msgList.add(userMsg);

        String prompt="角色:你是一个教务小助手;任务:请根据我们提供的课程信息，回答学生对课程安排的查询;例如:我说:软件3班,周一早上上什么课，你就回答:软件3班在周一早上上高等数学;举例完毕。" +
                "注意回答的要求:如果不能查询到准确信息，请回答，查询不到此信息，千万不要乱说。提供以下课程表资料供你参考：\n"+ragtxt.toString();
        System.out.println(prompt);
        Generation gen = new Generation();
        QwenParam params = QwenParam.builder().model(Generation.Models.QWEN_PLUS)
                //.apiKey(System.getenv("DASH_SCOPE_API_KEY"))//如果名字按规范来写，不需要设置，默认就是读这个名字
                //.prompt(prompt+"用户说:"+question)
                .messages(msgList)
                .seed(888)
                .topP(0.8)
                .resultFormat("message")
                .enableSearch(false)
                .maxTokens(maxTokens)
                .temperature((float)0.85)
                .repetitionPenalty((float)1.0)
                .build();

        GenerationResult result2 = gen.call(params);
        System.out.println(result2);
        for (GenerationOutput.Choice choice : result2.getOutput().getChoices()) {
            msgList.add(choice.getMessage());//大模型返回的消息记录到上下文
            // 如果需要调用工具
            if (result2.getOutput().getChoices().get(0).getMessage().getToolCalls() != null) {
                for (ToolCallBase toolCall : result2.getOutput().getChoices().get(0).getMessage()
                        .getToolCalls()) {//有可能要求助调用多个方法
                    if (toolCall.getType().equals("function")) {
                        // 获取工具函数名称和入参
                        String functionName = ((ToolCallFunction) toolCall).getFunction().getName();
                        String functionArgument = ((ToolCallFunction) toolCall).getFunction().getArguments();
                        // 大模型判断调用天气查询工具的情况
                        if (functionName.equals("get_current_week")) {
                            //帮LLM的忙，去调用自己的方法
                            Fun01.WeekTool weekTool=new Fun01.WeekTool();
                            String todayWeek=weekTool.weekToday();

                            Message toolResultMessage = Message.builder().role("tool")
                                    .content(String.valueOf(todayWeek)).toolCallId(toolCall.getId()).build();
                            msgList.add(toolResultMessage);
                            System.out.println("\n工具输出信息："+todayWeek);
                        }
                    }
                }
            }
            // 如果无需调用工具，直接输出大模型的回复
            else {
                System.out.println("\n最终答案："+result2.getOutput().getChoices().get(0).getMessage().getContent());
                return;
            }
        }
        // 大模型的第二轮调用 包含工具输出信息
        params.setMessages(msgList);
        result2 = gen.call(params);
        System.out.println("\n大模型第二轮输出信息："+JsonUtils.toJson(result));
        System.out.println(("\n最终答案："+result2.getOutput().getChoices().get(0).getMessage().getContent()));

    }
}

